package cn.spark.study.streaming;

import com.google.common.base.Optional;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.function.FlatMapFunction;
import org.apache.spark.api.java.function.Function2;
import org.apache.spark.api.java.function.PairFunction;
import org.apache.spark.streaming.Durations;
import org.apache.spark.streaming.api.java.JavaDStream;
import org.apache.spark.streaming.api.java.JavaPairDStream;
import org.apache.spark.streaming.api.java.JavaReceiverInputDStream;
import org.apache.spark.streaming.api.java.JavaStreamingContext;
import scala.Tuple2;

import java.util.Arrays;
import java.util.List;

/**
 * 基于updateStateByKey算子实现缓存机制的实时WordCount机制
 *
 * @author zhangj
 * @date 2020/11/18
 */
public class UpdateStateByKeyWordCount {
	public static void main(String[] args) {
		SparkConf conf = new SparkConf().setMaster("local[2]").setAppName("UpdateStateByKeyWordCount");

		JavaStreamingContext jssc = new JavaStreamingContext(conf, Durations.seconds(5));

		//如果要使用UpdateStateByKey算子,就必须设置一个checkpoint目录,开启checkpoint机制
		//以便于在内存数据丢失的时候,可以从checkpoint中恢复数据

		//开启checkpoint机制,只要调用jssc的checkpoint()方法,设置一个HDFS目录即可
		jssc.checkpoint("hdfs://ymm1:9000/wordcount_checkpoint");

		JavaReceiverInputDStream<String> lines = jssc.socketTextStream("ymm1", 9999);

		JavaDStream<String> words = lines.flatMap(new FlatMapFunction<String, String>() {
			@Override
			public Iterable<String> call(String line) throws Exception {
				return Arrays.asList(line.split(" "));
			}

		});

		JavaPairDStream<String, Integer> pairs = words.mapToPair(
				new PairFunction<String, String, Integer>() {
					@Override
					public Tuple2<String, Integer> call(String word) throws Exception {
						return new Tuple2<String, Integer>(word, 1);
					}
				});

		//到这里就不一样了,之前直接就是pairs.reduceByKey
		//然后就可以得到每个时间段的batch对应的RDD,计算出来的单词计数
		//但是这里有个问题,你如果要统计每个单词的全局计数呢?
		//就是说,统计出来,从程序启动,到现在为止,一个单词出现的次数
		//updateStateByKey,就可以实现直接通过Spark维护一份每个单词的全局计数
		JavaPairDStream<String, Integer> wordCounts = pairs.updateStateByKey(new Function2<List<Integer>, Optional<Integer>, Optional<Integer>>() {
			//对于每个单词,每次batch计算的时候,都会调用这个函数
			//values:相当于这个batch中,这个key的新值,可能有多个,比如说一个hello,可能有2个1,(hello,1),(hello,1),那么传入的是(1,1)
			//state:指这个key之前的状态,泛型是你自己指定的
			@Override
			public Optional<Integer> call(List<Integer> values, Optional<Integer> state) throws Exception {
				//首先定义一个全局的计数器
				Integer newValue = 0;

				//其次判断state是否存在,如果不存在说明这个Key是第一次出现
				//如果存在说明这个已经统计过全局的次数了
				if (state.isPresent()) {
					newValue = state.get();
				}

				//接着将本次新出现的值,都累加到newValue上,就是一个key目前的全局统计
				for (Integer value : values) {
					newValue += value;
				}
				return Optional.of(newValue);
			}
		});

		wordCounts.print();

		jssc.start();
		jssc.awaitTermination();
		jssc.close();
	}
}
